from sklearn.decomposition import PCA
from sklearn.preprocessing import RobustScaler

class DefenseEnhancer:
    def __init__(self, n_components=50):
        self.pca = PCA(n_components=n_components)
        self.scaler = RobustScaler()
        
    def adversarial_training(self, model, X_train, y_train, X_adv):
        """对抗训练"""
        combined_X = np.vstack([X_train, X_adv])
        combined_y = np.concatenate([y_train, np.ones(len(X_adv))])
        model.fit(combined_X, combined_y)
        return model
    
    def feature_squeezing(self, X):
        """特征压缩防御"""
        X_scaled = self.scaler.fit_transform(X)
        return self.pca.transform(X_scaled) 